# _*_ coding: utf-8 _*_
import pickle
from random import uniform

import numpy as np
from tqdm import trange

from Encodings import MultiToursSolution

import matplotlib.pyplot as plt

plt.figure(figsize=(12, 3))


def visualize(optimizer):
    plt.clf()
    plt.subplot(1, 3, 1)
    optimizer.data.draw(plt, routes=optimizer.best.tours, dot_size=0)
    plt.subplot(1, 3, 2)
    plt.imshow(optimizer.Pheromone)
    plt.title(f'{np.min(optimizer.Pheromone):.2g}, {np.max(optimizer.Pheromone):.2g}')
    plt.colorbar()
    plt.subplot(1, 3, 3)
    plt.plot(range(len(optimizer.evolve_list)), optimizer.evolve_list, label=f'{optimizer.best.total:.2f}')
    plt.legend()
    plt.pause(0.001)


def copy(target):
    return pickle.loads(pickle.dumps(target))


class ACO:  # Ant Colony Optimization
    def __init__(self, data):
        self.name = __class__.__name__
        # -- read instance from instance data --
        self.data = copy(data)
        self.V = self.data.cities  # city list
        self.n = self.data.n  # number of cities
        self.m = self.data.m  # number of salesmen
        self.A = self.data.A  # accessibility matrix
        self.Depot = self.data.depot  # depot
        self.V.remove(self.Depot)  # exclude depot
        self.Distance = self.data.distance_matrix  # distance matrix
        self.tour_cost = self.data.cost  # cost function for a single tour
        self.access = self.data.access  # fast computation access func
        # -- parameters ---
        self.NP = 10  # number of teams
        self.alpha = 1.0  # pheromone weight
        self.beta = 2.0  # heuristic weight
        self.rho = 0.1
        self.xee = 0.1
        self.q_zero = 0.9
        # -- solution initialization --
        self.optima = MultiToursSolution().init_solution(cities=self.V,
                                                         distance=self.Distance, access=self.access,
                                                         depot=self.Depot, m=self.m)
        self.tau_0 = 1.0 / (self.n * self.optima.total)
        self.Pheromone = np.full(shape=(self.n + 1, self.n + 1), fill_value=self.tau_0)
        self.evo_list = []
        # -- auxiliary --
        self.entropy = []  # assignment entropy to measure the diversity of population
    
    def optimize(self, iteration=1000, trial=1, diffuse=True):
        pbar = trange(iteration)
        pbar.set_description(f'{self.name} on {self.data.name} trial {trial}')
        for _ in pbar:
            # construct colony
            colony = [self.build_solution() for _ in range(self.NP)]  # build solutions
            pBest = min(colony, key=lambda _: _.total)
            # update optima & record evolutionary history
            if self.optima.total > pBest.total:  # update optima
                self.optima = copy(pBest)
            self.evo_list.append(self.optima.total)
            pbar.set_postfix_str(f'best: {self.optima.total:.2f}')
            # update pheromone
            self.update_pheromone()  # update pheromone
            if diffuse and self.is_stagnated():
                print(f'diffuse at {_}')
                self.diffuse_pheromone()  # diffuse pheromone
            
            visualize(self)
        plt.show()
        return self
    
    def build_solution(self):
        s = MultiToursSolution()
        s.tours = [[self.Depot] for _ in range(self.m)]
        intentions = np.zeros(self.m, dtype=int)
        cities = copy(self.V)
        while cities:
            intentions[:] = 0
            # collect intentions
            for k in range(self.m):
                curr = s.tours[k][-1]
                candidates = set(cities) & set(self.access(salesman=k))
                intentions[k] = self.transition(curr, candidates, q_zero=self.q_zero)
            # collision elimination
            intentions = self.collision_elimination(s, intentions)
            # ants step forward
            for k in range(self.m):
                if intentions[k] != 0:
                    s.tours[k].append(intentions[k])
                    # local pheromone update
                    a, b = s.tours[k][-1], s.tours[k][-2]
                    self.Pheromone[a][b] = (1 - self.xee) * self.Pheromone[a][b] + self.xee * self.tau_0
                    self.Pheromone[b][a] = (1 - self.xee) * self.Pheromone[b][a] + self.xee * self.tau_0
                    cities.remove(intentions[k])
        for k in range(self.m):
            s.tours[k].append(self.Depot)
        return s.calc_cost(cost_func=self.tour_cost)
    
    def collision_elimination(self, solution, intentions):
        for i in range(self.m):
            for j in range(i + 1, self.m):
                if intentions[i] == intentions[j]:
                    cost_i = self.Distance[solution.tours[i][-1]][intentions[i]]
                    cost_j = self.Distance[solution.tours[j][-1]][intentions[j]]
                    if cost_i > cost_j:
                        intentions[i] = 0
                    else:
                        intentions[j] = 0
        return intentions
    
    def transition(self, curr, roll, q_zero):
        """ Transition rule of ant colony system
        @param q_zero: criterion for exploitation
        @param curr: the current city where the ant is at
        @param roll: the roll list of cities accessible to the ant
        @return: the next city to reach
        """
        if len(roll) == 0:
            return 0
        roll = list(roll)
        prob = np.power(self.Pheromone[curr][roll], self.alpha) * np.power(1 / self.Distance[curr][roll], self.beta)
        prob = prob / prob.sum()
        if uniform(0, 1) < q_zero:
            city = roll[int(np.argmax(prob))]
        else:
            city = np.random.choice(a=roll, size=1, p=prob)[0]
        return city
    
    def update_pheromone(self):
        delta = 1.0 / self.optima.total
        for tour in self.optima.tours:
            for idx in range(1, len(tour) - 1):
                a, b = tour[idx - 1], tour[idx]
                self.Pheromone[a][b] = (1 - self.rho) * self.Pheromone[a][b] + self.rho * delta
                self.Pheromone[b][a] = (1 - self.rho) * self.Pheromone[b][a] + self.rho * delta
    
    def diffuse_pheromone(self):
        for tour in self.optima.tours:
            for idx in range(1, len(tour) - 1):
                a, b = tour[idx - 1], tour[idx]
                proximity = np.argsort(self.Distance[b])[2]
                self.Pheromone[proximity][b] += self.rho * self.Pheromone[a][b]
                self.Pheromone[b][proximity] += self.rho * self.Pheromone[b][a]
                self.Pheromone[a][b] = (1 - self.rho) * self.Pheromone[a][b]
                self.Pheromone[b][a] = (1 - self.rho) * self.Pheromone[b][a]
    
    def is_stagnated(self, criteria=100):
        return (len(self.evo_list) - self.evo_list.index(min(self.evo_list))) > criteria


from VCMTSP import Instance

case, m = 'att48', 2
data = Instance(file=f'../../data/TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)

aco_diffusion = ACO(data=data)
aco_diffusion.optimize(iteration=1000, diffuse=False)
quit(1)


results = []
for q_zero in [0.5, 0.6, 0.7, 0.8, 0.9]:
    aco_no_diffusion = ACO(data=data)
    aco_no_diffusion.q_zero = q_zero
    aco_no_diffusion.optimize(iteration=500, diffuse=False)
    results.append(aco_no_diffusion)


def visualize_results(results: list):
    if type(results) is not list:
        results = [results]
    plt.figure(figsize=(8, 6))
    plt.clf()
    for result in results:
        plt.plot(range(len(result.evo_list)),
                 result.evo_list,
                 label=f'{result.name}-q0={result.q_zero}:'
                       f' {result.optima.total:.2f}')
    plt.legend()
    plt.title(f'{results[0].data.name}: ')
    plt.show()


visualize_results(results)
